import cv2
import scipy.io as sio
import torch
import numpy as np

def batch_project(P, pts3d, K, angle_axis=True):
    n = pts3d.size(0)
    bs = P.size(0)
    device = P.device
    pts3d_h = torch.cat((pts3d, torch.ones(n, 1, device=device)), dim=-1)
    if angle_axis:
        R_out = kn.angle_axis_to_rotation_matrix(P[:, 0:3].view(bs, 3))
        PM = torch.cat((R_out[:,0:3,0:3], P[:, 3:6].view(bs, 3, 1)), dim=-1)
    else:
        PM = P
    pts3d_cam = pts3d_h.matmul(PM.transpose(-2,-1))
    pts2d_proj = pts3d_cam.matmul(K.t())
    S = pts2d_proj[:,:, 2].view(bs, n, 1)
    S[S==0] = S[S==0] + 1e-12
    pts2d_pro = pts2d_proj[:,:,0:2].div(S)

    return pts2d_pro

def compute_kp3d(poses,kp2d):
    corners = []
    for ind in range(6):
        A = np.zeros((len(kp2d)*2,3))
        p_is = [p[ind] for p in kp2d]
        p_norm = [np.linalg.inv(K)@np.array([p[0],p[1],1]) for p in p_is]
        for i in range(0,len(A),2):
            A[i] = p_norm[i//2][0] * poses[i//2][0][2] - poses[i//2][0][0]
            A[i+1] = p_norm[i//2][1] * poses[i//2][0][2] - poses[i//2][0][1]


        b = np.zeros((len(kp2d)*2,1))
        for i in range(0, len(A), 2):
            b[i] = poses[i//2][1][0] - p_norm[i//2][0]*poses[i//2][1][2]
            b[i+1] = poses[i//2][1][1] - p_norm[i//2][1]*poses[i//2][1][2]
        # b[0] = t1[0] - p1_norm[0]*t1[2]
        # b[1] = t1[1] - p1_norm[1]*t1[2]
        # b[2] = t2[0] - p2_norm[0]*t1[2]
        # b[3] = t2[1] - p2_norm[1]*t1[2]

        P,_,_,_ = np.linalg.lstsq(A,b,rcond=None)
        P = P.reshape(3)
        corners.append(P)
    return np.array(corners)


def get_K():
    fx = 24440.6
    fy = 24440.6
    u = 1280 / 2
    v = 720 / 2
    K = np.array(
        [[fx, 0, u],
         [0, fy, v],
         [0, 0, 1]],
    )
    return K


# 读取.mat 文件
K = get_K()

if __name__ == "__main__":
    from glob import glob
    import matplotlib.pyplot as plt

    kp2d = []
    poses = []
    path = "/media/liyuke/share/Micros_air2land/airport1_crash/002/left/"
    img_files = glob(path + '*color.jpg')
    mat_files = img_filenames = list(
        map(lambda x: x.strip().replace('color', 'meta').replace('.jpg', '.mat'), img_files))
    for i, file in enumerate(mat_files):
        data = sio.loadmat(file)
        kp2d.append(np.array(data["key_points"]))
        poses.append([data["relative_poses"][:3, :3], data["relative_poses"][:, 3]])
    corners = compute_kp3d(poses, kp2d)
    output_file = "/media/liyuke/share/Micros_air2land/3D_param.txt"
    np.savetxt(output_file, corners, fmt='%f')
    # test
    tets_path = "/media/liyuke/share/Micros_air2land/airport1_crosswind/001/left/"
    test_imgs = glob(tets_path + '*color.jpg')
    test_mats = img_filenames = list(
        map(lambda x: x.strip().replace('color', 'meta').replace('.jpg', '.mat'), test_imgs))

    for i, file in enumerate(test_mats):
        data = sio.loadmat(file)
        # kp2d.append(np.array(data["key_points"]))
        PM = torch.tensor(data["relative_poses"]).unsqueeze(0)
        pts2d = batch_project(PM, torch.tensor(corners).float(), torch.tensor(K).float(), angle_axis=False).numpy()[0]
        img = cv2.imread(test_imgs[i])
        for point in pts2d:
            cv2.circle(img, np.array([int(point[0]), int(point[1])]), 5, (0, 0, 255), -1)
        plt.figure()
        plt.imshow(img)
